/*
 *  jail_tty: spawn a program in its private controlling terminal.
 *  Copyright (C) 2011  David Decotigny <david.decotigny/at/google.com>
 *
 *  This program is free software: you can redistribute it and/or modify
 *  it under the terms of the GNU Affero General Public License as
 *  published by the Free Software Foundation, either version 3 of the
 *  License, or (at your option) any later version.
 *
 *  This program is distributed in the hope that it will be useful,
 *  but WITHOUT ANY WARRANTY; without even the implied warranty of
 *  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 *  GNU Affero General Public License for more details.
 *
 *  You should have received a copy of the GNU Affero General Public License
 *  along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

/*
 * Tool to spawn a program inside a private ctty (controlling tty):
 *   - creates a private controlling pty for child
 *   - configure a timeout alarm (if requested)
 *   - wait for the child to terminate or be timeout-ed and dump its stdout
 *   - return exit status of child (normal completion) or 255 (timeout)
 *
 * Useful for programs that hard-code user interaction through
 * /dev/tty (eg. programs that ask for a password, etc.).
 *
 * Tested:
 *   Ubuntu Lucid x86_64 (eglibc 2.10.1) / Linux 2.6.38
 *
 * Build:
 *   $ ./gcc -o jail_tty jail_tty.c -lutil
 *
 * Examples:
 *   # Plain execution:
 *   $ ./jail_tty -- sh -c "echo HELLO ; sleep 5 ; echo BYE ; exit 42"
 *   HELLO
 *   BYE
 *   # (retval = 42)
 *
 *   # Execution aborted by timeout:
 *   $ ./jail_tty -a 2 -- sh -c "echo HELLO ; sleep 5 ; echo BYE ; exit 42"
 *   HELLO
 *   # (retval = 255 [timeout])
 */
#include <pty.h>
#include <stdlib.h>
#include <unistd.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <stdio.h>
#include <assert.h>
#include <getopt.h>
#include <alloca.h>


/*
 * Print program usage info to stderr.
 *
 * Args:
 *   argv0: name/path of program
 *
 * Returns:
 *  -1
 */
static int usage(const char* argv0) {
  fprintf(stderr, "Usage: %s [options] -- child_prog child_args...\n", argv0);
  fprintf(stderr, "\n");
  fprintf(stderr, "Options:\n");
  fprintf(stderr, "  -a|--alarm seconds   Fire a timeout (seconds > 0)\n");
  fprintf(stderr, "  -h|--help            This cruft\n");
  return -1;
}


/*
 * Run child program with args in a private CTTY.
 *
 * Child stdout/stderr will be printed to jail_tty's stdout. Nothing
 * can be sent to child's stdin.
 *
 * Args:
 *   timeout_secs: seconds before firing an alarm, <= 0 to wait forever
 *   child_argc: number of entries in child_argv
 *   child_argv: array of child parameter strings
 *
 * Returns:
 *   -1 on error or timeout, child exit status otherwise.
 */
static int do_jail_pty_with_timeout(int timeout_secs,
                                    int child_argc,
                                    char * const child_argv[]) {
  char ** exec_argv;
  int i;
  int amaster;
  int fork_res;

  assert(child_argc >= 1);

  exec_argv = alloca(sizeof(char*) * (child_argc+1));
  for (i = 0 ; i < child_argc ; ++i)
    exec_argv[i] = child_argv[i];
  exec_argv[i] = NULL;

  fork_res = forkpty(& amaster, NULL, NULL, NULL);
  if (fork_res < 0) {
    perror("fork_pty");
    return -1;
  }

  /* In child */
  if (fork_res == 0) {
    if (timeout_secs > 0)
      alarm(timeout_secs);
    else
      alarm(0); /* cancel any pending alarm */

    execvp(exec_argv[0], exec_argv);
    perror("execvp");
    return -1;
  }

  /* In father */
  {
    char buff[4096];
    int child_status = 0;

    while (1) {
      int read_res = read(amaster, buff, sizeof(buff));
      if (read_res <= 0)
        break;
      if (read_res != write(fileno(stdout), buff, read_res))
        perror("write");
    }
    if (fork_res != waitpid(fork_res, & child_status, 0)) {
      perror("waitpid");
      return -1;
    }
    if (WIFEXITED(child_status))
      return WEXITSTATUS(child_status);
    return -1;
  }
}


int main(int argc, char * const argv[]) {
  int timeout_secs = -1;

  while (1) {
    int c;
    static struct option long_options[] = {
      {"help", 0, 0, 'h'},
      {"alarm", 1, 0, 'a'},
      {0, 0, 0, 0}
    };

    c = getopt_long(argc, argv, "ha:", long_options, NULL);
    if (c < 0)
      break;
    switch(c) {
      case 'h':
        {
          usage(argv[0]);
          return 0;
        }

      case 'a':
        {
          if (sscanf(optarg, "%i", & timeout_secs) != 1)
            return usage(argv[0]);
          if (timeout_secs <= 0)
            return usage(argv[0]);
          break;
        }

      default:
        return usage(argv[0]);
    }
  }

  if (optind >= argc)
    return usage(argv[0]);

  return do_jail_pty_with_timeout(timeout_secs, argc-optind, &argv[optind]);
}
